import requests
import tinify
import sys
import os
from colorama import init, Fore

# 初始化颜色输出
init(autoreset=True)

# 多个 Tinify API Key
TINIFY_KEYS = [
    "JThxRRFtSKT1qjnjqm0gNgHm4yJfNXn0",
    "7mN2HCh8w1rf90VwmjFSnZSXDqK96cMh",
    # 可以在这里添加更多 API Key
    # "your_second_key",
    # "your_third_key",
]

current_key_index = 0
failed_keys = []

def switch_to_next_key():
    global current_key_index
    current_key_index += 1
    if current_key_index < len(TINIFY_KEYS):
        tinify.key = TINIFY_KEYS[current_key_index]
        print(Fore.YELLOW + f'🔁 切换到下一个 API Key: {TINIFY_KEYS[current_key_index]}')
        return True
    else:
        print(Fore.RED + '❌ 所有 Tinify API Key 已耗尽或失效！')
        return False

def setup_first_valid_key():
    global current_key_index
    for index, key in enumerate(TINIFY_KEYS):
        tinify.key = key
        try:
            # 尝试获取用量，判断是否有效
            _ = tinify.compression_count
            current_key_index = index
            print(Fore.GREEN + f'✅ 当前使用 API Key: {TINIFY_KEYS[current_key_index]}')
            return True
        except Exception as e:
            failed_keys.append(key)
    return False

def compress_image(image_path):
    global current_key_index
    while current_key_index < len(TINIFY_KEYS):
        try:
            print(Fore.YELLOW + f'[START] Compressing: {image_path}')
            source = tinify.from_file(image_path)
            source.to_file(image_path)
            print(Fore.GREEN + f'[OK]     Compressed: {image_path}\n')
            return True
        except tinify.AccountError:
            print(Fore.RED + f"[LIMIT] 当前 API Key 已达到压缩次数限制：{TINIFY_KEYS[current_key_index]}")
            failed_keys.append(TINIFY_KEYS[current_key_index])
            if not switch_to_next_key():
                break
        except Exception as e:
            print(Fore.RED + f"[ERROR] 压缩失败: {e}")
            return False
    return False

def compress_images_in_folder(folder_path):
    count = 0
    success = 0
    for root, dirs, files in os.walk(folder_path):
        for file in files:
            if file.lower().endswith(('.png', '.jpg')):
                image_path = os.path.join(root, file)
                if compress_image(image_path):
                    success += 1
                count += 1
    print(Fore.CYAN + f'📊 本文件夹共检测到 {count} 张图片，成功压缩 {success} 张。\n')

def show_quota():
    print(Fore.MAGENTA + "\n📦 各 API Key 使用情况：")
    TEST_IMAGE = "https://tinypng.com/images/panda-happy.png"
    for i, key in enumerate(TINIFY_KEYS):
        tinify.key = key
        try:
            # 为了获得真实的额度主动触发一次压缩（不保存）
            response = requests.get(TEST_IMAGE)
            tinify.from_buffer(response.content).to_buffer()
            
            usage = tinify.compression_count
            print(Fore.MAGENTA + f'  🔑 Key {i + 1}: 已用 {usage} / 500 次')
        except Exception:
            print(Fore.RED + f'  🔑 Key {i + 1}: 获取额度失败 (可能无效或超额)')

def main():
    if len(sys.argv) < 2:
        print(Fore.RED + "用法: python compressor.py <图片路径/文件夹路径> [更多路径...]")
        sys.exit(1)

    if not setup_first_valid_key():
        print(Fore.RED + "❌ 所有 Tinify API Key 都不可用，程序终止。")
        sys.exit(1)

    for target_path in sys.argv[1:]:
        print(Fore.BLUE + f'\n🚀 正在处理: {target_path}\n')

        if not os.path.exists(target_path):
            print(Fore.RED + f"[ERROR] 路径不存在: {target_path}")
            continue

        if os.path.isdir(target_path):
            compress_images_in_folder(target_path)
        else:
            compress_image(target_path)

    show_quota()
    input(Fore.MAGENTA + "🎉 所有任务完成，按任意键退出。")

if __name__ == "__main__":
    main()
